import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, PoseWithCovarianceStamped, TwistWithCovariance, TransformStamped,Quaternion
from sensor_msgs.msg import NavSatFix,NavSatStatus
from nav_msgs.msg import Odometry,Path
from tf2_ros import TransformBroadcaster
import numpy as np

class ROS2Driver(Node):
    def __init__(self):
        super().__init__('rtk_node')
        
        # 声明参数
        self.DeclareParams()
        # 开启接口
        self.DeclareInterface()
        
    def DeclareInterface(self):
        '''申明接口'''
        self.pose_pub = self.create_publisher(PoseWithCovarianceStamped, 'pos_rtk', 10)
        self.odom_pub = self.create_publisher(Odometry, 'odom_rtk', 10)
        self.path_pub = self.create_publisher(Path, 'path_rtk', 10)
        self.lla_pub = self.create_publisher(NavSatFix, 'lla_rtk', 10)
        self.path_poses = []  # 用于累积路径点
        #   TF发布器
        self.tf_broadcaster = TransformBroadcaster(self)
        
    def DeclareParams(self):
        '''参数申明'''
        self.declare_parameter('data_type', "wifi")
        self.declare_parameter('dev_id', "ZDYDGPS250006")
        self.declare_parameter('uart_port', "/dev/ttyUSB0")
        self.declare_parameter('http_port', 8000)
        self.declare_parameter('buffer_len', 100)
        self.declare_parameter('publish_odom', True)
        self.declare_parameter('publish_path', True)
        self.declare_parameter('publish_tf', True)
        self.declare_parameter('publish_lla', True)
        self.declare_parameter('frame_id', "rtk_gps")
        self.declare_parameter('global_frame', "map")
        self.data_type=self.get_parameter('data_type').get_parameter_value().string_value
        self.dev_id=self.get_parameter('dev_id').get_parameter_value().string_value
        self.uart_port=self.get_parameter('uart_port').get_parameter_value().string_value
        self.buffer_len=self.get_parameter("buffer_len").get_parameter_value().integer_value
        self.publish_odom = self.get_parameter('publish_odom').get_parameter_value().bool_value
        self.publish_tf= self.get_parameter('publish_tf').get_parameter_value().bool_value
        self.publish_path= self.get_parameter('publish_path').get_parameter_value().bool_value
        self.publish_lla= self.get_parameter('publish_lla').get_parameter_value().bool_value
        self.frame_id=self.get_parameter('frame_id').get_parameter_value().string_value
        self.global_frame=self.get_parameter('global_frame').get_parameter_value().string_value
        self.http_port=self.get_parameter("http_port").get_parameter_value().integer_value

    def Publish(self, x, y, z,lon,lat,height, roll=0.0,pitch=0.0,yaw=0.0,speed=0.0,rms=0.0,status=''):
        """
            x (float): x坐标
            y (float): y坐标
            z (float): z坐标
            yaw (float): 水平朝向角度(弧度)
            rms (float): 均方根误差值
        """
        current_time = self.get_clock().now().to_msg()
        states={'Non':-1,'NONE':-1,'None':-1,'GPS':0,'RTK*':0,'RTK':1}
        # 基于RMS值计算协方差矩阵（RMS越大，协方差越大）
        pose_covariance = self._calculate_covariance(rms)
        twist_covariance = self._calculate_covariance(rms, is_twist=True)
        
        # 发布PoseWithCovarianceStamped消息
        pose_msg = PoseWithCovarianceStamped()
        pose_msg.header.stamp = current_time
        pose_msg.header.frame_id = 'map'
        pose_msg.pose.pose.position.x = x
        pose_msg.pose.pose.position.y = y
        pose_msg.pose.pose.position.z = z
        # 转换欧拉角为四元数
        q = self._euler_to_quaternion(roll, pitch, yaw)
        pose_msg.pose.pose.orientation=q
        
        # 设置协方差矩阵
        pose_msg.pose.covariance = pose_covariance
        self.pose_pub.publish(pose_msg)
        
        # 根据参数决定是否发布Odometry消息
        if self.publish_odom:
            # 发布Odometry消息
            odom_msg = Odometry()
            odom_msg.header.stamp = current_time
            odom_msg.header.frame_id = self.global_frame
            odom_msg.child_frame_id = self.frame_id
            # 设置位姿
            odom_msg.pose.pose.position.x = x
            odom_msg.pose.pose.position.y = y
            odom_msg.pose.pose.position.z = z
            odom_msg.pose.pose.orientation=q
            # 设置位姿协方差
            odom_msg.pose.covariance = pose_covariance
            # 设置速度(当前仅支持x轴速度)
            odom_msg.twist.twist.linear.x =speed
            odom_msg.twist.twist.linear.y = 0.0
            odom_msg.twist.twist.linear.z = 0.0
            odom_msg.twist.twist.angular.x = 0.0
            odom_msg.twist.twist.angular.y = 0.0
            odom_msg.twist.twist.angular.z = 0.0
            # 设置速度协方差
            odom_msg.twist.covariance = twist_covariance
            self.odom_pub.publish(odom_msg)
        if self.publish_tf:
            #发布坐标系变换
            tf_msg = TransformStamped()
            tf_msg.header.stamp = current_time
            tf_msg.header.frame_id = self.global_frame
            tf_msg.child_frame_id = self.frame_id
            
            tf_msg.transform.translation.x = x
            tf_msg.transform.translation.y = y
            tf_msg.transform.translation.z = z
            
            tf_msg.transform.rotation=q
            
            self.tf_broadcaster.sendTransform(tf_msg)
        if self.publish_path:
            #发布路径
            path_msg = Path()
            path_msg.header.stamp = current_time
            path_msg.header.frame_id = self.global_frame  # 路径基于全局坐标系
            
            # 创建当前位姿的PoseStamped
            current_pose = PoseStamped()
            current_pose.header.stamp = current_time
            current_pose.header.frame_id = self.global_frame
            current_pose.pose.position.x = x
            current_pose.pose.position.y = y
            current_pose.pose.position.z = z
            current_pose.pose.orientation = q  # 与当前位姿一致的朝向
            self.path_poses.append(current_pose)
            if len(self.path_poses) > self.buffer_len: self.path_poses.pop(0)  # 移除最早的点
            path_msg.poses = self.path_poses
            self.path_pub.publish(path_msg)
        if self.publish_lla:
            # 发布LLA坐标
            lla_data=NavSatFix()
            lla_data.header.stamp = current_time
            lla_data.header.frame_id=self.frame_id
            lla_data.latitude=lat
            lla_data.longitude=lon
            lla_data.altitude=height
            lla_data.position_covariance=[rms, 0.0, 0.0,
                                        0.0, rms, 0.0,
                                        0.0 ,0.0 ,rms]
            if status in states.keys(): lla_data.status.status=states[status]
            else: lla_data.status.status=-1
            self.lla_pub.publish(lla_data)
        
    def _euler_to_quaternion(self, roll, pitch, yaw):
        """
        参数:
            roll (float): 绕x轴旋转角度(弧度)
            pitch (float): 绕y轴旋转角度(弧度)
            yaw (float): 绕z轴旋转角度(弧度)
        
        返回:
            Quaternion: 四元数消息
        """
        # 创建四元数消息
        q = Quaternion()
        
        # 四元数计算
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)
        
        q.w = cr * cp * cy + sr * sp * sy
        q.x = sr * cp * cy - cr * sp * sy
        q.y = cr * sp * cy + sr * cp * sy
        q.z = cr * cp * sy - sr * sp * cy
        
        return q
    def _calculate_covariance(self, rms, is_twist=False):
        """
        基于RMS值计算协方差矩阵
        
        参数:
            rms (float): 均方根误差值
            is_twist (bool): 是否为速度协方差矩阵
        
        返回:
            list: 扁平化的协方差矩阵
        """
        # 基础协方差值
        base_covariance = np.diag([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
        
        # 速度协方差通常较小
        if is_twist:
            base_covariance = np.diag([0.01, 0.01, 0.01, 0.01, 0.01, 0.01])
        
        # 协方差与RMS的平方成正比（使不确定性增长更显著）
        covariance = base_covariance * (rms ** 2)
        
        # 确保协方差不为零
        covariance = np.maximum(covariance, 1e-6 * np.eye(6))
        
        return covariance.flatten().tolist()

